package org.whale.controller;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.HuggingFaceTokenizer;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.*;
import org.whale.assistant.DoctorAgent;
import org.whale.dto.ChatMessageDTO;
import reactor.core.publisher.Flux;

import java.util.List;

import static org.springframework.http.MediaType.TEXT_EVENT_STREAM_VALUE;

@RestController
public class DoctorController {

    @Resource
    private DoctorAgent doctorAgent;
    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;
    @Resource
    private EmbeddingModel embeddingModel;

    @PostMapping(value = "/stream/chat",produces = "text/stream;charset=utf-8")
    public Flux<String> streamChat(@RequestBody ChatMessageDTO chatMessageDTO){
        try{
            return doctorAgent.streamChat(chatMessageDTO.getMessage(),chatMessageDTO.getMemoryId());
        }catch (Exception e){
            e.printStackTrace();
            return Flux.just("康康有点忙不过来啦，请稍后重新提问");
        }
    }


    @GetMapping("/rag/load")
    public String reagLoad(){
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("C:\\Users\\Administrator\\Desktop\\rag");
        EmbeddingStoreIngestor.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                //文档切割 - 按照什么样的规则把文档分段
                .documentSplitter(new DocumentByParagraphSplitter(60,10,new HuggingFaceTokenizer()))
                .build().ingest(documents);

        return "success";
    }


}
